#include "ScreenShot.h"

#include "stb_image_write.h"
#include "wlr-screencopy.h"

#include <fcntl.h>
#include <sys/mman.h>
#include <unistd.h>

wayland::display_t display;
wayland::registry_t registry;
wayland::shm_t shm;
wayland::output_t output;
wayland::zwlr_screencopy_manager_v1_t screencopyManager;

wayland::shm_pool_t pool;
std::vector<wayland::output_t> data;

bool buffer_copy_done = false;

static struct
{
    uint8_t *data;
    struct wayland::buffer_t wl_buffer;
    enum wayland::shm_format format;
    int width, height, stride;
    bool y_invert;
} buffer;

// Add definition of your processing function here
void ScreenShot::getScreenshotImg(const HttpRequestPtr &req, std::function<void(const HttpResponsePtr &)> &&callback)
{
    if (!registry) {
        registry = display.get_registry();
        registry.on_global() = [](uint32_t name, const std::string &interface, uint32_t version) {
            if (interface == wayland::shm_t::interface_name) {
                shm = static_cast<wayland::shm_t>(registry.bind(name, shm, std::min<int>(wayland::detail::shm_interface.version, version)));
            } else if (interface == wayland::output_t::interface_name) {
                output = static_cast<wayland::output_t>(registry.bind(name, output, std::min<int>(wayland::detail::output_interface.version, version)));
            } else if (interface == wayland::zwlr_screencopy_manager_v1_t::interface_name) {
                screencopyManager = static_cast<wayland::zwlr_screencopy_manager_v1_t>(registry.bind(name, screencopyManager, std::min<int>(wayland::detail::zwlr_screencopy_manager_v1_interface.version, version)));
            }
        };
        display.roundtrip();
    }

    buffer_copy_done = false;
    wayland::zwlr_screencopy_frame_v1_t screencopyFrame = screencopyManager.capture_output(0, output);
    screencopyFrame.on_ready() = [&](uint32_t, uint32_t, uint32_t) {
        buffer_copy_done = true;
    };
    screencopyFrame.on_flags() = [&](wayland::zwlr_screencopy_frame_v1_flags flags) {
        buffer.y_invert = flags & wayland::zwlr_screencopy_frame_v1_flags::y_invert;
    };
    screencopyFrame.on_failed() = []() {
        std::cerr << "Failed to copy frame" << std::endl;
        exit(1);
    };
    screencopyFrame.on_buffer() = [&](wayland::shm_format format, uint32_t width, uint32_t height, uint32_t stride) {
        auto old_format = buffer.format;
        buffer.format = static_cast<wayland::shm_format>(format);
        buffer.width = width;
        buffer.height = height;
        buffer.stride = stride;

        if (format != old_format || !buffer.wl_buffer) {
            if (buffer.data) {
                munmap(buffer.data, buffer.stride * buffer.height);
            }

            std::string tmp_file = "/tmp/screencopy-XXXXXX";
            int fd = mkstemp((char *)tmp_file.c_str());
            if (fd < 0) {
                return;
            }

            int ret;
            while ((ret = ftruncate(fd, stride * height)) == EINTR) {
                // No-op
            }
            if (ret < 0) {
                close(fd);
                return;
            }
            unlink(tmp_file.c_str());

            buffer.data = (uint8_t *)mmap(NULL, stride * height, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
            if (buffer.data == MAP_FAILED) {
                fprintf(stderr, "mmap failed: %m\n");
                close(fd);
                return;
            }

            wayland::shm_pool_t pool = shm.create_pool(fd, stride * height);
            close(fd);
            buffer.wl_buffer = pool.create_buffer(0, width, height, stride, format);
        }

        screencopyFrame.copy(buffer.wl_buffer);
    };

    while (!buffer_copy_done && display.dispatch() != -1) {
        // No-op
    }

    if (stbi_write_png("/tmp/drogon-screencopy.png", buffer.width, buffer.height, buffer.stride / buffer.width, buffer.data, buffer.stride) == 0) {
        auto resp = HttpResponse::newHttpResponse();
        resp->setBody("Failed to write PNG");
        callback(resp);
        return;
    }

    // return the screenshot as a response
    auto httpFileResponse = HttpResponse::newFileResponse("/tmp/drogon-screencopy.png");
    httpFileResponse->setContentTypeCode(CT_APPLICATION_OCTET_STREAM);
    httpFileResponse->setStatusCode(k200OK);
    callback(httpFileResponse);
}

void ScreenShot::getScreenshot(const HttpRequestPtr &req, std::function<void(const HttpResponsePtr &)> &&callback)
{
    auto resp = HttpResponse::newHttpResponse();
    resp->setBody("<image src=\"/screenshot.png\" alt=\"Screenshot\"/>");
    callback(resp);
}